from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.colors import TwoSlopeNorm

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)


def get_intrinsic_pc( params ):

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    # """ load model
    # """
    # params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        rnn_state_trajectory.append(np.array(rnn_state).copy())
        
        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 对 rnn_state_np 进行 PCA
    pca = PCA()
    # pca.fit(rnn_state_np)
    pca.fit(rnn_state_trajectory_np)

    # 打印 variance ratio
    print(pca.explained_variance_ratio_)

    rnn_state_np_pca = pca.transform(rnn_state_np)

    # 创建KMeans对象，指定聚类数为4
    kmeans = KMeans(n_clusters=4)
    # 对rnn_state_np_pca进行聚类
    kmeans.fit(rnn_state_np_pca)
    # 获取聚类中心的坐标
    cluster_centers = kmeans.cluster_centers_

    return pca, cluster_centers


def constraint_run( params , dims = 50):

    pca, cluster_centers = get_intrinsic_pc( params )

    pcs = pca.components_.copy()

    pca_mean = pca.mean_.copy()

    print("shape of pcs: ", pcs.shape)
    print("shape of pca_mean: ", pca_mean.shape)

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    # """ load model
    # """
    # params = load_weights(rpl_config.model_pth)
    # # 定义一个函数，用于生成随机权重
    # def init_weights_r(key, shape):
    #     return jax.random.normal(key, shape)
    # # 生成一个随机的 PRNGKey
    # key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # # key = jax.random.PRNGKey(4512)
    # # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    # random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    # params = random_params

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])

    mat_intr = np.array(tree_leaves[1])[0:128,:]
    print("mat_intr.shape: ", mat_intr.shape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 1000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)

    print("rnn_state.shape: ", rnn_state.shape)

    rnn_state_trajectory = []
    intr_field = []

    pc_id = 23
    pc_ = pcs[pc_id]

    # 生成一个包含随机个数的整数集合，元素个数大于20个
    num_indices = np.random.randint(21, 128)
    # 生成一个包含[0, 127]范围内所有整数的数组
    all_indices = np.arange(128)
    # 从中随机选择num_indices个不重复的元素
    random_indices = np.random.choice(all_indices, size=num_indices, replace=False)
    random_indices.sort()

    # selected_pcs = [0,1,2,3,4]
    # selected_pcs = [i for i in range(dims)]
    selected_pcs = [i for i in random_indices]
    pcs_selected = pcs[selected_pcs]

    print("selected_pcs :",selected_pcs)

    # 生成一个和 pcs_selected 形状一样的，元素范围在(-1,1) 区间的随机向量
    pcs_rnd = np.random.uniform(-1, 1, pcs_selected.shape)
    pcs_rnd_norm = np.linalg.norm(pcs_rnd, axis=1)
    pcs_rnd = pcs_rnd / pcs_rnd_norm[:,None]

    # # 应用随机主成分向量
    pcs_selected = pcs_rnd

    rnn_state_old = rnn_state.copy()

    var_log = []

    for t in range(rpl_config.life_duration):

        # progress_bar(t, rpl_config.life_duration)

        intr_vector = np.dot(rnn_state, mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field.append(intr_vector - rnn_state)
        rnn_state_trajectory.append(np.array(rnn_state).copy())
        # rnn_state = intr_vector

        # # 使用单一的维度重投影
        # dot_prod = np.dot(intr_vector-pca_mean, pc_)
        # rnn_state = np.outer(dot_prod, pc_) + pca_mean

        # 使用多个维度重投影
        dot_prods = np.dot(intr_vector-pca_mean, pcs_selected.T)
        # 将dot_prods转置为(128, 1000)形状的矩阵
        dot_prods_T = np.transpose(dot_prods)
        # 将dot_prods_T与pcs相乘，得到还原后的向量
        restored_vectors_T = np.dot(pcs_selected.T, dot_prods_T)
        # 将还原后的向量转置为(1000, 128)形状的矩阵
        restored_vectors = np.transpose(restored_vectors_T) + pca_mean
        diff_norm = np.linalg.norm(restored_vectors - intr_vector, axis=1)
        # print("diff_norm: ", diff_norm)
        rnn_state = restored_vectors
        
        rnn_state_diff = rnn_state - rnn_state_old
        rnn_state_diff_norm = np.linalg.norm(rnn_state_diff, axis=1)
        sum_norm = np.sum(rnn_state_diff_norm)
        # print("sum_norm: ", sum_norm)

        var_log.append(sum_norm)

        rnn_state_old = rnn_state.copy()

    print("sum_norm: ", sum_norm)

    # 将 var_log 绘制成曲线
    var_log = np.array(var_log)
    plt.plot(var_log)
    plt.show()

    # # 将 pcs_selected 中的没一个向量，绘制成 bar chart
    # for i in range(pcs_selected.shape[0]):
    #     plt.subplot(1, pcs_selected.shape[0], i+1)
    #     plt.bar(range(pcs_selected.shape[1]), pcs_selected[i])
    # plt.show()
    
if __name__ == "__main__":


    # test_vec1 = np.array([[1,2,1],[2,3,2],[1,2,1],[2,3,2]])
    # test_vec2 = np.array([[3,3,3]])
    # test_vec2_norm = np.linalg.norm(test_vec2, axis=1)
    # # test_vec2 = test_vec2 / test_vec2_norm[:,None]
    # print(test_vec2)
    # print("shape of test_vec1: ", test_vec1.shape)
    # print("shape of test_vec2: ", test_vec2.shape)
    # res = np.dot(test_vec1, test_vec2.T)
    # print(res)
    # rep = np.dot(test_vec2.T, res.T)
    # print(rep)

    # exit()

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    # 定义一个函数，用于生成随机权重
    def init_weights_r(key, shape):
        return jax.random.normal(key, shape)
    # 生成一个随机的 PRNGKey
    key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # key = jax.random.PRNGKey(4512)
    # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    
    params = random_params
    
    constraint_run( params , 60)